#coding:utf-8
import os, cv2, argparse
import numpy as np
from model import DetHeight
import torch
from ipdb import set_trace
from torch.autograd import Variable
import torch.backends.cudnn as cudnn

parser = argparse.ArgumentParser(description='Get height of an image')
parser.add_argument('--model_path', default='/checkpoint/model.pth',
                    help='path to model')
parser.add_argument('--image_path', default='test.jpg',
                    help='path to test image')
parser.add_argument('--gpu', default=-1, type=int, 
    help='gpu device, -1 to use cpu, 0 to use the first gpu')
args = parser.parse_args()
if args.gpu != -1:
    torch.cuda.set_device(args.gpu)
    cudnn.benchmark = True

def readimg(imgpath, scale_size=(640, 360)):
    img = cv2.imread(os.path.join(imgpath))
    if img.shape[0] != scale_size[1] or img.shape[1] != scale_size[0]:
        img = cv2.resize(img, scale_size)
    img = np.array(img, dtype=np.float32)
    img = img / 255.0 - 0.5
    img = np.transpose(img, axes=[2,0,1])
    return img

def get_height(model_path, image_path, gpu):
    if os.path.exists(model_path) and os.path.exists(image_path):
        img = readimg(imgpath=image_path)
        img = Variable(torch.from_numpy(img[None,:,:,:]))
        if gpu != -1:
            img = img.cuda()
        model = DetHeight()
        if gpu != -1:
            model = model.cuda()
        model.load_state_dict(torch.load(model_path))
        model.eval()
        output = model(img)
        if args.gpu != -1:
            output = output.data.cpu().numpy()
        else:
            output = output.data.numpy()
        return output
    return 1000

if __name__ == '__main__':
    height = get_height(model_path=args.model_path, image_path=args.image_path, gpu=args.gpu)
    print('{"z": %.2f, "x": %.2f, "y":  %.2f}' % (height, 1000, 1000 ))

